from io import BytesIO
from PIL import Image
from nicegui import ui
import base64
import imghdr
import json
import requests

from utils import binary_to_base64, generate_random_ipv4_address, calculate_md5


def get_img_b64(file_obj):
    img_format = imghdr.what(file_obj)
    if not img_format:
        return False, "图片格式错误"
    img = Image.open(file_obj)
    width, height = img.size
    ratio = max(width / 2500, height / 2500)
    img = img.resize((int(width / ratio), int(height / ratio)))
    buffer = BytesIO()
    img.save(buffer, img_format)
    base64_string = binary_to_base64(buffer.getvalue())
    return True, f'data:image/{img_format.lower()};base64,{base64_string}'


def get_img_md5(b64_str: str):
    data = json.dumps({"photo": b64_str})
    md5_bytes = calculate_md5(data)
    return binary_to_base64(md5_bytes)


def remove_bg(file_obj, app=1):
    ok, img_b64_str = get_img_b64(file_obj)
    if not ok:
        return False, img_b64_str
    img_md5_str = get_img_md5(img_b64_str)

    ip = generate_random_ipv4_address()
    headers = {
        'Content-Type': 'application/json',
        'X-Forwarded-For': ip,
        'X-Real-IP': ip
    }
    params = {
        "content_md5": img_md5_str,
        "app": app
    }
    resp = requests.get("https://api-quzuotu.segapi.com/photo/iapi/request_auth", params=params, headers=headers)
    auth_resp = resp.json()
    if auth_resp.get("status") != 0:
        return False, auth_resp
    auth_header = resp.json()['data']["headers"]

    url = "https://iapi.aisegment.com/segment/portrait/rgba" \
        if app == 1 else "https://iapi.aisegment.com/segment/common/rgba"

    headers = {
        'Content-Type': 'application/json',
        'X-Forwarded-For': ip,
        'X-Real-IP': ip,
        **auth_header
    }
    resp = requests.post(url, json={"photo": img_b64_str}, headers=headers).json()
    if resp.get("status") == 0:
        return True, resp["data"]
    return False, resp.get("errmsg")
